import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

# 设置字体和样式
plt.rcParams['font.family'] = 'DejaVu Sans'
plt.rcParams['axes.unicode_minus'] = False
plt.rcParams['axes.grid'] = True
plt.rcParams['grid.alpha'] = 0.3
plt.rcParams['grid.linestyle'] = '--'

# 读取数据
data_fix = pd.read_csv('data-10fix.csv', index_col=0)
data_v1g = pd.read_csv('data-10v1g.csv', index_col=0)
data_v2g = pd.read_csv('data-10v2g.csv', index_col=0)

# 设置颜色
colors = {
    'Nuclear': '#FF6B6B',      # 鲜红色
    'Wind': '#4ECDC4',         # 青绿色
    'Central_PV': '#FFD93D',   # 明黄色
    'Coal': '#95A5A6',         # 灰色
    'Gas': '#2ECC71',         # 绿色
    'Hydro_Nonpumped': '#3498DB'  # 蓝色
}

# 创建图形
fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(20, 8))

# 绘制三个子图
def plot_stacked_bar(data, ax, title):
    bottom = np.zeros(len(data.columns))
    for source in data.index:
        ax.bar(range(len(data.columns)), data.loc[source], bottom=bottom, 
               label=source, color=colors[source], width=0.8)
        bottom += data.loc[source]
    
    # 设置x轴标签
    ax.set_xticks(range(len(data.columns)))
    years = [year[12:] for year in data.columns]
    ax.set_xticklabels(years, rotation=45)
    
    # 设置标题和标签
    ax.set_title(title, pad=20, fontsize=12, fontweight='bold')
    ax.set_ylabel('Power Generation (TWh)', fontsize=10)
    
    # 设置y轴范围和刻度
    ax.set_ylim(0, 12000)
    ax.yaxis.set_major_locator(plt.MultipleLocator(2000))

# 绘制三个子图
plot_stacked_bar(data_fix, ax1, 'Base Scenario')
plot_stacked_bar(data_v1g, ax2, 'Scenario 1')
plot_stacked_bar(data_v2g, ax3, 'Scenario 2')

# 添加总标题
fig.suptitle('Energy Structure Changes in Different Scenarios', fontsize=14, fontweight='bold', y=1.02)

# 添加图例
handles, labels = ax1.get_legend_handles_labels()
fig.legend(handles, labels, loc='center right', bbox_to_anchor=(1.15, 0.5),
          title='Energy Types', title_fontsize=10, fontsize=9)

# 调整布局
plt.tight_layout()

# 保存图片
plt.savefig('energy_structure_comparison.png', dpi=300, bbox_inches='tight')
plt.close() 